package com.citusdata.elven.httpRpc;

import com.citusdata.elven.httpRpc.HttpRpcUrl;
import com.citusdata.elven.util.ConfigManager;
import com.google.protobuf.Descriptors.MethodDescriptor;
import com.google.protobuf.Descriptors.ServiceDescriptor;
import com.google.protobuf.InvalidProtocolBufferException;
import com.google.protobuf.UninitializedMessageException;
import com.google.protobuf.Message;
import com.google.protobuf.RpcCallback;
import com.google.protobuf.RpcChannel;
import com.google.protobuf.RpcController;
import java.io.IOException;
import java.net.Socket;
import java.net.InetSocketAddress;
import java.net.UnknownHostException;

import org.apache.commons.configuration.Configuration;
import org.apache.http.*;
import org.apache.http.entity.*;
import org.apache.http.impl.*;
import org.apache.http.message.*;
import org.apache.http.params.*;
import org.apache.http.protocol.*;
import org.apache.http.util.*;
import org.apache.log4j.Logger;


public class HttpRpcChannel implements RpcChannel {
  private static final Logger logger = Logger.getLogger(HttpRpcChannel.class);
  private static final Configuration config = ConfigManager.getConfig();
  private static final byte[] NULL_RESPONSE = new byte[0];

  private final String  host;
  private final int     port;
  private final boolean persistent;
  private final int     connTimeout;

  private final HttpParams httpParams = new BasicHttpParams();
  private final DefaultHttpClientConnection httpConnection 
    = new DefaultHttpClientConnection();

  public HttpRpcChannel(String host, int port, boolean persistent) {
    this.host = host;
    this.port = port;
    this.persistent  = persistent;
    this.connTimeout = config.getInt("HttpRpcChannel.ConnectionTimeout", 2*1000);

    // Look up timeout configuration from HttpRpcServer.SocketTimeout (default:
    // 10 seconds) so that both the client and server ends of a connection get a
    // timeout at around the same time.
    int socketTimeout = config.getInt("HttpRpcServer.SocketTimeout", 10*1000);
    httpParams.setIntParameter(CoreConnectionPNames.SO_TIMEOUT, socketTimeout);
  }

  protected boolean connected() {
    // Check whether the connection is still active. Note that this check will
    // not detect connections closed after some time of inactivity, and the next
    // operation on the socket will throw an IOException in those cases. Using
    // a relatively expensive "staleness" check (10-30ms) will solve this
    // problem.
    return httpConnection.isOpen();
  }

  protected void shutdown() {
    try {
      if (connected()) {
        // Most likely, this connection was created in a different thread.
        // Therefore, shutdown() is the only safe way to close this connection. 
        httpConnection.shutdown();
      }
    } catch (IOException ignore) { }
  }


  private boolean connect(HttpRpcController httpController) {
    Socket socket = null;

    try {
      HttpProtocolParams.setVersion(httpParams, HttpVersion.HTTP_1_1);

      // HttpParams do not set the connection timeout, therefore we manually set
      // it here. We first create an empty socket and then connect using the
      // preconfigured timeout value.
      socket = new Socket();
      socket.connect(new InetSocketAddress(host, port), connTimeout);

      httpConnection.bind(socket, httpParams);
    } catch (UnknownHostException ex) {
      handleError(httpController, "Unknown host", HttpStatus.SC_NOT_FOUND, ex);
    } catch (IOException ex) {
      handleError(httpController, "I/O error", HttpStatus.SC_NOT_FOUND, ex);
      try {
        if (socket != null) {
          socket.close();
        }
      } catch (IOException ignore) { }
    }

    boolean connected = connected();
    return connected;
  }


  private boolean reuseConnection(HttpResponse response, HttpContext context) {
    if (persistent) {
      ConnectionReuseStrategy reuseStrategy = new DefaultConnectionReuseStrategy();
      if (reuseStrategy.keepAlive(response, context)) {
        return true;
      }
    } 

    return false;
  }


  private byte[] dispatchRequest(String requestUrl, byte[] requestData,
                                 HttpRpcController httpController) {

    HttpHost endpoint = new HttpHost(host, port);
    HttpContext context = new BasicHttpContext();
    context.setAttribute(ExecutionContext.HTTP_TARGET_HOST, endpoint);
    context.setAttribute(ExecutionContext.HTTP_CONNECTION, httpConnection);

    BasicHttpProcessor processor = new BasicHttpProcessor();
    processor.addInterceptor(new RequestContent());
    processor.addInterceptor(new RequestTargetHost());
    processor.addInterceptor(new RequestConnControl());

    BasicHttpEntityEnclosingRequest request
      = new BasicHttpEntityEnclosingRequest(HttpRpcUrl.METHOD, requestUrl);
    request.setEntity(new ByteArrayEntity(requestData));
    request.setParams(httpParams);

    // Initialize the responseData to a zero-length array
    byte[] responseData = NULL_RESPONSE;

    try {
      HttpRequestExecutor executor = new HttpRequestExecutor();
      executor.preProcess(request, processor, context);
    
      HttpResponse response = executor.execute(request, httpConnection, context);    
      
      response.setParams(httpParams);
      executor.postProcess(response, processor, context);

      // Process response status and data returned by the server
      StatusLine responseStatus = response.getStatusLine();
      if (responseStatus.getStatusCode() == HttpStatus.SC_OK) {
        HttpEntity entity = response.getEntity();
        responseData = EntityUtils.toByteArray(entity);
      } else {
        handleError(httpController, responseStatus.getReasonPhrase(), 
                    responseStatus.getStatusCode(), null);
      }

      // If the connection can not be reused, close connection.
      if (!reuseConnection(response, context)) {
        shutdown();
      }

    } catch (IOException ex) {
      handleError(httpController, "I/O error", HttpStatus.SC_NOT_FOUND, ex);
      shutdown();
    } catch (HttpException ex) {
      handleError(httpController, "Protocol error", HttpStatus.SC_NOT_FOUND, ex);
      shutdown();
    }

    return responseData;
  }


  public void callMethod(MethodDescriptor method, RpcController controller,
                         Message request, Message responsePrototype, 
                         RpcCallback done) {
    // The caller must pass in an HttpRpcController.
    HttpRpcController httpController = (HttpRpcController) controller;
    
    if (!request.isInitialized()) {
      handleError(httpController, "Malformed request object", 
                  HttpStatus.SC_BAD_REQUEST, null);
      return;
    }

    // Connect to the host only if we dont have an open connection to the host.
    if (!connected()) {
      boolean established = connect(httpController);
      if (!established) {
        // handleError() already called by the previous function
        return;
      }
    }

    // Create and dispatch the Http request.
    ServiceDescriptor service = method.getService();
    HttpRpcUrl requestUrl = new HttpRpcUrl(service.getName(), method.getName());

    byte[] response = dispatchRequest(requestUrl.getUrl(), request.toByteArray(),
                                      httpController);
    if (httpController.failed()) {
      // handleError() already called by the previous function
      return;
    }

    // Create response prototype from byte array, and call registered callbacks.
    try {
      Message.Builder builder = responsePrototype.newBuilderForType();
      builder.mergeFrom(response);
      responsePrototype = builder.build();
      if (done != null) {
        done.run(responsePrototype);
      }
    } catch (InvalidProtocolBufferException ex) {
      handleError(httpController, "Invalid response protocol",
                  HttpStatus.SC_UNPROCESSABLE_ENTITY, ex);
    } catch (UninitializedMessageException ex) {
      handleError(httpController, "Response protocol missing required fields",
                  HttpStatus.SC_UNPROCESSABLE_ENTITY, ex);
    }
  }


  public void dispatchUrl(String requestUrl) {
    // This convenience function sends simple urls to the server
    byte[] emptyRequest = NULL_RESPONSE;
    HttpRpcController httpController = new HttpRpcController();

    if (!connected()) {
      boolean established = connect(httpController);
      if (!established) {
        // handleError() already called by connect()
        return;
      }
    }

    dispatchRequest(requestUrl, emptyRequest, httpController);
  }


  private void handleError(HttpRpcController httpController,
                           String message, int statusCode, Exception ex) {
    logger.warn(message + ": " + statusCode + " failure with exception: " + ex);
    httpController.setFailed(message, statusCode);
  }
}